package com.gitee.wsl.flow.coroutines

import com.gitee.wsl.ext.base.currentTimeMillis
import com.gitee.wsl.collections.list.LinkedList
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import kotlinx.coroutines.yield
import kotlin.time.Duration


fun <T> Flow<T>.rateLimit(limit: Long, per: Duration): Flow<T> {
    return rateLimit(limit, per.inWholeMilliseconds)
}

fun <T> Flow<T>.rateLimit(limit: Long, per: Long) = flow<T> {
    coroutineScope {
        val context = coroutineContext
        var lastStartTime = currentTimeMillis
        var remainingInDuration = limit
        val items = LinkedList<T>()
        var isDone = false
        launch(context) {
            collect {
                items.add(it)
            }
            isDone = true
        }
        launch(Dispatchers.Default) {
            while (isActive) {
                yield()
                if (remainingInDuration > 0) {
                    val removeFirst = items.removeFirstOrNull()
                    if (removeFirst != null) {
                        withContext(context) {
                            emit(removeFirst)
                        }
                        remainingInDuration--
                    } else {
                        if (isDone) {
                            break
                        }
                    }

                } else {
                    val waitUntil = lastStartTime + per
                    delay(waitUntil - currentTimeMillis)
                    lastStartTime = currentTimeMillis
                    remainingInDuration = limit
                }
            }
        }
    }
}